import torch  # 导入PyTorch库
import torch.nn as nn  # 导入PyTorch的神经网络模块
import torch.optim as optim  # 导入PyTorch的优化器模块
import torch.nn.functional as F  # 导入PyTorch的函数式模块
from torch.utils.data import DataLoader, TensorDataset  # 导入PyTorch的数据加载工具


# # 定义transformer模型
# class Transformer(nn.Module):  # 定义一个继承自nn.Module的Transformer类
#     def __init__(self, input_dim, output_dim, hidden_dim, num_layers):  # 初始化函数，定义模型的结构和参数
#         super(Transformer, self).__init__()  # 调用父类的初始化函数
#         self.encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(input_dim, nhead=2),
#                                              num_layers)  # 定义Transformer编码器层
#         self.fc = nn.Linear(input_dim, output_dim)  # 定义线性层
#
#     def forward(self, x):  # 定义模型的前向传播过程
#         x = self.encoder(x)  # 经过编码器层
#         x = self.fc(x)  # 经过线性层
#         return x  # 返回输出结果
#
#
# # 构建数据集
# input_dim = 10  # 输入维度
# output_dim = 1  # 输出维度
# hidden_dim = 16  # 隐藏层维度
# num_layers = 2  # Transformer层数
# batch_size = 32  # 批量大小
# num_samples = 1000  # 样本数量
#
# X = torch.randn(num_samples, input_dim)  # 随机生成输入数据
# y = torch.randn(num_samples, output_dim)  # 随机生成输出数据
#
# dataset = TensorDataset(X, y)  # 构建PyTorch的数据集
# dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)  # 构建数据加载器，用于批量加载数据
#
# # 初始化模型和优化器
# model = Transformer(input_dim, output_dim, hidden_dim, num_layers)  # 初始化Transformer模型
# criterion = nn.MSELoss()  # 定义损失函数
# optimizer = optim.Adam(model.parameters(), lr=0.01)  # 定义优化器，使用Adam算法进行优化
#
# # 训练模型
# num_epochs = 10  # 训练轮数
#
# for epoch in range(num_epochs):  # 循环训练轮数
#     for inputs, targets in dataloader:  # 遍历每个批次的数据
#         optimizer.zero_grad()  # 梯度清零
#         outputs = model(inputs)  # 前向传播，得到模型输出
#         loss = criterion(outputs, targets)  # 计算损失
#         loss.backward()  # 反向传播，计算梯度
#         optimizer.step()  # 更新模型参数
#
#     print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')  # 打印当前训练轮数和损失值
#
# # 使用模型进行预测
# new_inputs = torch.randn(1, input_dim)  # 随机生成新的输入数据
# predicted_output = model(new_inputs)  # 使用模型进行预测
# print("Predicted output:", predicted_output)  # 打印预测结果

# 创建一个包含10个词的词嵌入，每个词嵌入的维度为5
# embedding_layer = nn.Embedding(num_embeddings=10, embedding_dim=5)
# layer = embedding_layer(torch.tensor([1, 2, 3, 4]))
# print(layer)


import torch
#
# def check_torch_gpu():
#     if torch.cuda.is_available():
#         print(torch.cuda.get_device_name(0))
#     else:
#         print("No GPU available.")
#
# check_torch_gpu()

import torch
import torchvision

device = torch.device("mps")

x = torch.randn(128, 128, device=device)
net = torchvision.models.resnet18().to(device)

print(x.device)
print(next(net.parameters()).device)

if __name__ == '__main__':
    #embedding_layer = nn.Embedding(num_embeddings=10, embedding_dim=5)
    print("over")

# 预测下一个词